import torch
from torch import nn
from torch.nn import functional as F
from torchvision import models


# 定义特征提取器（原 extractors.py 的内容）
class ResNet34(nn.Module):
    def __init__(self, ):
        super().__init__()
        resnet = models.resnet34()

        # 提取 ResNet34 的前几层
        self.initial = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
        )
        self.layer1 = resnet.layer1  # 64 channels
        self.layer2 = resnet.layer2  # 128 channels
        self.layer3 = resnet.layer3  # 256 channels
        self.layer4 = resnet.layer4  # 512 channels

    def forward(self, x):
        x = self.initial(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        class_f = x  # 用于辅助分类器的特征
        x = self.layer4(x)
        return x, class_f  # 返回主分支特征和辅助分类特征


def resnet34():
    return ResNet34()


# 主模型代码
class PSPModule(nn.Module):
    def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):
        super().__init__()
        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
        self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)
        self.relu = nn.ReLU()

    def _make_stage(self, features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [F.interpolate(stage(feats), size=(h, w), mode='bilinear', align_corners=True) for stage in
                  self.stages]
        priors.append(feats)
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)


class PSPUpsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.PReLU()
        )

    def forward(self, x):
        h, w = 2 * x.size(2), 2 * x.size(3)
        p = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)
        return self.conv(p)


class PSPNet(nn.Module):
    def __init__(self, n_classes=18, sizes=(1, 2, 3, 6), backend='resnet34'):
        super().__init__()

        # 根据 backend 选择特征提取器
        if backend == 'resnet34':
            self.feats = resnet34()
            psp_size = 512
            deep_features_size = 256
        else:
            raise ValueError('Unsupported backend: {}'.format(backend))

        self.psp = PSPModule(psp_size, 1024, sizes)
        self.drop_1 = nn.Dropout2d(p=0.3)

        self.up_1 = PSPUpsample(1024, 256)
        self.up_2 = PSPUpsample(256, 64)
        self.up_3 = PSPUpsample(64, 64)

        self.drop_2 = nn.Dropout2d(p=0.15)
        self.final = nn.Sequential(
            nn.Conv2d(64, n_classes, kernel_size=1),
            nn.LogSoftmax(dim=1)
        )

        self.classifier = nn.Sequential(
            nn.Linear(deep_features_size, 256),
            nn.ReLU(),
            nn.Linear(256, n_classes)
        )

    def forward(self, x):
        f, class_f = self.feats(x)
        p = self.psp(f)
        p = self.drop_1(p)

        p = self.up_1(p)
        p = self.drop_2(p)

        p = self.up_2(p)
        p = self.drop_2(p)

        p = self.up_3(p)
        p = self.drop_2(p)

        auxiliary = F.adaptive_max_pool2d(class_f, output_size=(1, 1)).view(-1, class_f.size(1))
        return self.final(p), self.classifier(auxiliary)


# 示例用法
if __name__ == '__main__':
    model = PSPNet(n_classes=18, backend='resnet34')
    input_tensor = torch.randn(1, 3, 224, 224)
    output, auxiliary_output = model(input_tensor)
    print('主输出形状:', output.shape)  # 例如: [1, 18, 56, 56]
    print('辅助分类器输出形状:', auxiliary_output.shape)  # 例如: [1, 18]
